import os
import logging
from flask import redirect, flash, render_template, url_for
# 从werkzeug.utils导入secure_filename函数，用于安全处理文件名
from werkzeug.utils import secure_filename
# 从utils模块导入allowed_image_file函数，用于判断文件类型是否合法
from utils import allowed_image_file
# 导入腾讯云COS上传工具
from cos_utils import upload_to_cos
from embedding_utils import get_image_embedding
# 从milvus_utils模块导入插入向量的函数
from milvus_utils import insert_image_vectors, insert_image_vector, search_image_vectors

logger = logging.getLogger(__name__)


def handle_image_upload(request):
    if request.method == 'POST':
        if "file" not in request.files:
            logger.warning("文件名为空")
            flash("文件名为空", "error")
            return redirect(request.url)
        file = request.files["file"]
        # 检查文件名是否为空
        if file.filename == "":
            # 记录文件名为空的日志
            logger.warning("未选择文件")
            # 弹出提示信息
            flash("未选择文件", "error")
            # 重定向回当前页面
            return redirect(request.url)

        # 判断文件存在且为允许的图片类型
        if file and allowed_image_file(file.filename):
            # 获取安全的文件名
            filename = secure_filename(file.filename)
            try:
                # 上传图片到腾讯云COS，返回图片url
                image_url = upload_to_cos(file.stream, key=filename)
                logger.info(f"图片上传成功，URL：{image_url}")
                #获取图片的向量表示
                embedding = get_image_embedding(image_url)
                logger.info("图片向量化成功，开始入库")
                #向Milvus插入图片向量和图片url
                insert_image_vector(embedding, image_url)
                logger.info("图片向量入库成功")
                # 闪现成功提示信息
                flash("文件上传并入库成功", "success")
            except Exception as e:
                logger.error(f"向量化或入库失败{e}")
                # 闪现失败提示信息
                flash(f"文件上传或入库失败{e}", "error")

            return redirect(url_for("upload_image"))
    # 渲染图片上传页面
    return render_template("upload_image.html")

def handle_image_search(request):
    # 初始化检索结果列表
    results = []
    if request.method == 'POST':
        if "file" not in request.files:
            logger.warning("文件名为空")
            flash("文件名为空", "error")
            return redirect(request.url)
        file = request.files["file"]
        # 检查文件名是否为空
        if file.filename == "":
            # 记录文件名为空的日志
            logger.warning("未选择文件")
            # 弹出提示信息
            flash("未选择文件", "error")
            # 重定向回当前页面
            return redirect(request.url)

        # 判断文件存在且为允许的图片类型
        if file and allowed_image_file(file.filename):
            # 获取安全的文件名
            filename = secure_filename(file.filename)
            try:
                # 上传图片到腾讯云COS，返回图片url
                image_url = upload_to_cos(file.stream, key=filename)
                logger.info(f"图片上传成功，URL：{image_url}")
                # 获取图片的向量表示
                query_embedding = get_image_embedding(image_url)
                logger.info("图片向量化成功，开始Milvus检索")
                milvus_results = search_image_vectors(query_embedding, top_k = 5)
                logger.info(f"Milvus图片检索完成， 命中组数：{len(milvus_results)}")
                # 遍历检索结果，提取图片URL和相似度分数
                for hits in milvus_results:
                    for hit in hits:
                        img_url = hit.entity.get("image_url", "")
                        results.append({
                            "image_url": img_url,
                            "score": hit.distance
                        })
            except Exception as e:
                # 记录一场日职并闪现错误细腻下
                logger.error(f"图片检索失败{e}")
                flash(f"图片检索失败{e}", "error")

    # 渲染图片检索结果页面
    return render_template("search_image.html", results=results)
